import copy
import math

import common_modules

from DrawingAPIs.line_drawer import *

# Parse the overall loss value trend. Return x axis labels and loss values.
def OverallLoss(file):
    labels = []
    loss_vals = []
    with open(file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            labels.append(row[0].split('-')[1])
            loss_vals.append(float(row[3]))
    return labels, loss_vals

# Parse separate loss values. {output name: [loss values]}
def SepLoss(file):
    loss_vals = {}
    with open(file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            if row[0] not in loss_vals.keys():
                loss_vals.setdefault(row[0], [])
            loss_vals[row[0]].append(float(row[2]))
    return loss_vals

def DrawConvergeLossLine(args):
    original_overall_loss_file = args.root[0] + "/CIFAR10LossK40c/Original/CIFAR10_nvOverallLoss.csv"
    original_sep_loss_file = args.root[0] + "/CIFAR10LossK40c/Original/CIFAR10_nvSepLoss.csv"
    intra_overall_loss_file = args.root[0] + "/CIFAR10LossK40c/IntraLayer/CIFAR10_nvOverallLoss.csv"
    intra_sep_loss_file = args.root[0] + "/CIFAR10LossK40c/IntraLayer/CIFAR10_nvSepLoss.csv"
    pipe_overall_loss_file = args.root[0] + "/CIFAR10LossK40c/Pipeline/CIFAR10_nvOverallLoss.csv"
    pipe_sep_loss_file = args.root[0] + "/CIFAR10LossK40c/Pipeline/CIFAR10_nvSepLoss.csv"
    # original_overall_loss_file = args.root[0] + "/GoogleNetLossK40c/Original/GoogleNetOverallLoss.csv"
    # original_sep_loss_file = args.root[0] + "/GoogleNetLossK40c/Original/GoogleNetSepLoss.csv"
    # intra_overall_loss_file = args.root[0] + "/GoogleNetLossK40c/IntraLayer/GoogleNetOverallLoss.csv"
    # intra_sep_loss_file = args.root[0] + "/GoogleNetLossK40c/IntraLayer/GoogleNetSepLoss.csv"

    x_labels, original_overall_loss_vals = OverallLoss(original_overall_loss_file)
    original_sep_loss_vals = SepLoss(original_sep_loss_file)
    unused_x_labels, intra_overall_loss_vals = OverallLoss(intra_overall_loss_file)
    intra_sep_loss_vals = SepLoss(intra_sep_loss_file)
    unused_pipe_x_labels, pipe_overall_loss_vals = OverallLoss(pipe_overall_loss_file)
    # pipe_sep_loss_vals = SepLoss(pipe_overall_loss_file)

    x_tick_set = []
    x_tick_set.append(np.arange(1, len(x_labels) + 1, 1))
    x_tick_set.append(np.arange(1, len(x_labels) + 1, 1))
    x_tick_set.append(np.arange(1, len(x_labels) + 1, 1))

    overall_y_max = max(max(original_overall_loss_vals), max(intra_overall_loss_vals))
    overall_y_max = max(overall_y_max, max(pipe_overall_loss_vals))

    if overall_y_max > math.ceil(overall_y_max) - 0.5:
        overall_y_max = math.ceil(overall_y_max)
    else:
        overall_y_max = math.ceil(overall_y_max)

    overall_val_set = [original_overall_loss_vals, intra_overall_loss_vals, pipe_overall_loss_vals]
    # sep_val_set = [original_sep_loss_vals, intra_sep_loss_vals, pipe_sep_loss_vals]
    legend_labels = ['Caffe', 'GLP4NN', 'HGP4NN']

    x_label_pos = [0]
    x_ticklabels = ['0']
    interval = len(x_labels) // 5
    for i in range(1, 6, 1):
        x_label_pos.append(int(x_labels[i * interval - 1]) / 100)
        x_ticklabels.append(x_labels[i * interval - 1])

    # Draw overall converge trend.
    print("Drawing OverallLoss.pdf...")
    DrawLineFig(x_tick_set, x_label_pos, x_ticklabels, overall_val_set, legend_labels,
                x_label='Iteration', y_label='Loss', x_min=0, x_max=(len(x_labels) + 1),
                y_min=0, y_max=overall_y_max, folder=(args.output[0] + '/Loss/'),
                filename='exp_conv_loss', allow_legend=True)
